
#include"undistort.h"

#include <pcl_conversions/pcl_conversions.h>
#include <pcl/common/transforms.h>


//根据比例来进行四元数插值，首先根据两个四元数来计算相对的四元数，然后进行插值
//template <typename PointType>
Eigen::Quaterniond getInterpolatedAttitude(const Eigen::Quaterniond &q_s_w, const Eigen::Quaterniond &q_e_w, float scale)
{
    if (0 == scale || scale > 1)
        return move(Eigen::Quaterniond().Identity());

    // calculate angleaxis difference
    Eigen::Quaterniond q_e_s = q_s_w.inverse() * q_e_w; //inv(start)*end，计算出相对旋转 
    q_e_s.normalize();
    Eigen::AngleAxisd diff_angle_axis(q_e_s);//将四元数转换为旋转轴和旋转角度

    // interpolated attitude by scale
    float interpolated_angle = diff_angle_axis.angle() * scale;//对旋转角度进行缩放
    Eigen::Quaterniond q_ie_s(Eigen::AngleAxisd(interpolated_angle, diff_angle_axis.axis()).toRotationMatrix());
    Eigen::Quaterniond q_ie_w = q_s_w * q_ie_s;
    q_ie_w.normalize();

    return move(q_ie_w);
}



CImuIntegration::CImuIntegration()
{

}
CImuIntegration::~CImuIntegration()
{

}

int CImuIntegration::AddImu(ImuData& imu)
{

    if(mvImuData.size()==0)
    {
        imu.rot.setIdentity();
        imu.rpy[0] = 0;
        imu.rpy[1] = 0;
        imu.rpy[2] = 0;
        mvImuData.push_back(imu);
    }
    else
    {
        int size = mvImuData.size();

        //标定代码里面的旋转积分算法，直接计算出表示旋转的四元数
        Eigen::Vector3d bar_gyr    = 0.5 * (mvImuData[size-1].gyr + imu.gyr);
        Eigen::Vector3d angle_inc  = bar_gyr * (imu.stamp - mvImuData[size-1].stamp);
        Eigen::Quaterniond rot_inc = Eigen::Quaterniond(1.0, 0.5 * angle_inc[0], 0.5 * angle_inc[1], 0.5 * angle_inc[2]);
        imu.rot = mvImuData[size-1].rot * rot_inc;
        //std::cout<<imu.rot.coeffs()<<std::endl;

        //lio-sam里面直接对三个角度进行累加
        double lastStamp = mvImuData[size-1].stamp;
        double curStamp = imu.stamp;
        double dt = curStamp - lastStamp;
        // imu.rpy[0] = mvImuData[size-1].rpy[0] + dt*imu.gyr[0];
        // imu.rpy[1] = mvImuData[size-1].rpy[1] + dt*imu.gyr[1];
        // imu.rpy[2] = mvImuData[size-1].rpy[2] + dt*imu.gyr[2];
        imu.rpy = mvImuData[size-1].rpy + dt*bar_gyr;
        Eigen::Affine3f rotation = pcl::getTransformation(0, 0, 0, imu.rpy[0], imu.rpy[1], imu.rpy[2]);
        Eigen::Matrix3f m = (rotation.matrix()).block<3,3>(0,0);
        Eigen::Quaternionf q(m);
        //std::cout<<q.coeffs()<<std::endl;
        //std::cout<<std::endl;
        //Eigen::Vector3d rpy = imu_buffer_[i].rot.matrix().eulerAngles(0, 1, 2);

        mvImuData.push_back(imu);

    }

    mvImuData.push_back(imu);

    return 0;
}


int CImuIntegration::FindAngle(int start, double time, double& ax, double& ay, double& az)
{
    int index = 0;

    for(int i=start; i<mvImuData.size(); i++)
    {
        if(time<mvImuData[i].stamp)
        {
            index = i;
            break;
        }
    }

    if(index>0)
    {
        int li = index-1;
        int ri = index;
        double lt = mvImuData[li].stamp;
        double rt = mvImuData[ri].stamp;
        double w = (time - lt) / (rt-lt);
        Eigen::Vector3d rpy = (1-w)*mvImuData[li].rpy + w*mvImuData[ri].rpy;
        ax = rpy[0];
        ay = rpy[1];
        az = rpy[2];
    }
   
    return index;
}


int CImuIntegration::InterpolateQuaternion(int start, double time, Eigen::Quaterniond& iq)
{
    int index = 0;

    for(int i=start; i<mvImuData.size(); i++)
    {
        if(time<mvImuData[i].stamp)
        {
            index = i;
            break;
        }
    }

    if(index>0)
    {
        int li = index-1;
        int ri = index;
        double lt = mvImuData[li].stamp;
        double rt = mvImuData[ri].stamp;
        double scale = (time - lt) / (rt-lt);

        Eigen::Quaterniond q_b1_w = mvImuData[li].rot;
        Eigen::Quaterniond q_b2_w = mvImuData[ri].rot;
        
        iq = getInterpolatedAttitude(q_b1_w, q_b2_w, scale);
    }
   
    return index;
}




CLidarUnDistort::CLidarUnDistort()
{

}
CLidarUnDistort::~CLidarUnDistort()
{

}

int CLidarUnDistort::AddImu(ImuData& imu)
{

    if(mvImuData.size()==0)
    {
        imu.rot.setIdentity();
        imu.rpy[0] = 0;
        imu.rpy[1] = 0;
        imu.rpy[2] = 0;
        mvImuData.push_back(imu);
    }
    else
    {
        int size = mvImuData.size();

        //标定代码里面的旋转积分算法，直接计算出表示旋转的四元数
        Eigen::Vector3d bar_gyr    = 0.5 * (mvImuData[size-1].gyr + imu.gyr);
        Eigen::Vector3d angle_inc  = bar_gyr * (imu.stamp - mvImuData[size-1].stamp);
        Eigen::Quaterniond rot_inc = Eigen::Quaterniond(1.0, 0.5 * angle_inc[0], 0.5 * angle_inc[1], 0.5 * angle_inc[2]);
        imu.rot = mvImuData[size-1].rot * rot_inc;
        //std::cout<<imu.rot.coeffs()<<std::endl;

        //lio-sam里面直接对三个角度进行累加
        double lastStamp = mvImuData[size-1].stamp;
        double curStamp = imu.stamp;
        double dt = curStamp - lastStamp;
        // imu.rpy[0] = mvImuData[size-1].rpy[0] + dt*imu.gyr[0];
        // imu.rpy[1] = mvImuData[size-1].rpy[1] + dt*imu.gyr[1];
        // imu.rpy[2] = mvImuData[size-1].rpy[2] + dt*imu.gyr[2];
        imu.rpy = mvImuData[size-1].rpy + dt*bar_gyr;
        Eigen::Affine3f rotation = pcl::getTransformation(0, 0, 0, imu.rpy[0], imu.rpy[1], imu.rpy[2]);
        Eigen::Matrix3f m = (rotation.matrix()).block<3,3>(0,0);
        Eigen::Quaternionf q(m);
        //std::cout<<q.coeffs()<<std::endl;
        //std::cout<<std::endl;
        //Eigen::Vector3d rpy = imu_buffer_[i].rot.matrix().eulerAngles(0, 1, 2);

        mvImuData.push_back(imu);

    }

    mvImuData.push_back(imu);

    return 0;
}


int CLidarUnDistort::FindAngle(int start, double time, double& ax, double& ay, double& az)
{
    int index = 0;

    for(int i=start; i<mvImuData.size(); i++)
    {
        if(time<mvImuData[i].stamp)
        {
            index = i;
            break;
        }
    }

    if(index>0)
    {
        int li = index-1;
        int ri = index;
        double lt = mvImuData[li].stamp;
        double rt = mvImuData[ri].stamp;
        double w = (time - lt) / (rt-lt);
        Eigen::Vector3d rpy = (1-w)*mvImuData[li].rpy + w*mvImuData[ri].rpy;
        ax = rpy[0];
        ay = rpy[1];
        az = rpy[2];
    }
   
    return index;
}


int CLidarUnDistort::InterpolateQuaternion(int start, double time, Eigen::Quaterniond& iq)
{
    int index = 0;

    for(int i=start; i<mvImuData.size(); i++)
    {
        if(time<mvImuData[i].stamp)
        {
            index = i;
            break;
        }
    }

    if(index>0)
    {
        int li = index-1;
        int ri = index;
        double lt = mvImuData[li].stamp;
        double rt = mvImuData[ri].stamp;
        double scale = (time - lt) / (rt-lt);

        Eigen::Quaterniond q_b1_w = mvImuData[li].rot;
        Eigen::Quaterniond q_b2_w = mvImuData[ri].rot;
        
        iq = getInterpolatedAttitude(q_b1_w, q_b2_w, scale);
    }
   
    return index;
}


Eigen::Affine3f CLidarUnDistort::PredictIncreRot(double beginStamp, double endStamp)
{
    Eigen::Affine3f increRot;
    increRot.setIdentity();

    if(endStamp>beginStamp)
    {
        double ax,ay,az;
        int imuIndex = FindAngle(0, beginStamp, ax, ay, az);
        if(imuIndex>0)   
        {
            Eigen::Affine3f beginRot = pcl::getTransformation(0, 0, 0, ax, ay, az);
            FindAngle(imuIndex, endStamp, ax, ay, az);   
            Eigen::Affine3f endRot = pcl::getTransformation(0, 0, 0, ax, ay, az);
            increRot = beginRot.inverse()*endRot;
        }
    }

    return increRot;
}

//根据IMU提供的角速度来进行旋转运动畸变校正
int CLidarUnDistort::Process(pcl::PointCloud<PointXYZIRT>& scan, double initStamp)
{
    //find the initial rotation angle
    //float lidarTime = scan.points[0].time;
    double ax,ay,az;
    int imuIndex = FindAngle(0, initStamp, ax, ay, az);

    //如果第一个点落在imu的有效时间之外，那么就返回-1
    if(imuIndex<=0)
    {
        return -1;
    }

    //the first point rotation
    Eigen::Affine3f initRot = pcl::getTransformation(0, 0, 0, ax, ay, az);
    
    for (int k = 0; k < scan.size(); k++)
    {
        //计算每个点对应的旋转角度
        double ptTime = initStamp + scan.points[k].time;
        imuIndex = FindAngle(imuIndex, ptTime, ax, ay, az);

        //计算当前点和第一个点之间的相对旋转矩阵
        Eigen::Affine3f curRot = pcl::getTransformation(0, 0, 0, ax, ay, az);
        Eigen::Affine3f transBt = initRot.inverse()*curRot; //relative rotation

        //根据相对旋转矩阵对点云进行运动畸变校正
        PointXYZIRT* point = &scan.points[k];
        float x = transBt(0,0) * point->x + transBt(0,1) * point->y + transBt(0,2) * point->z + transBt(0,3);
        float y = transBt(1,0) * point->x + transBt(1,1) * point->y + transBt(1,2) * point->z + transBt(1,3);
        float z = transBt(2,0) * point->x + transBt(2,1) * point->y + transBt(2,2) * point->z + transBt(2,3);
        point->x = x;
        point->y = y;
        point->z = z;
    }

    return 0;
}

